import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as transforms
import time
import sys

# 设置随机种子以确保结果可重现
torch.manual_seed(42)

# 检查设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 定义数据预处理
transform_train = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

# 加载CIFAR10数据集
def load_cifar10_data():
    """加载CIFAR10数据集"""
    print("正在加载CIFAR10数据集...")
    
    try:
        trainset = torchvision.datasets.CIFAR10(
            root='./data', train=True, download=True, transform=transform_train)
        trainloader = DataLoader(
            trainset, batch_size=128, shuffle=True, num_workers=0)  # num_workers=0避免多进程问题
        
        testset = torchvision.datasets.CIFAR10(
            root='./data', train=False, download=True, transform=transform_test)
        testloader = DataLoader(
            testset, batch_size=128, shuffle=False, num_workers=0)
        
        classes = ('plane', 'car', 'bird', 'cat', 'deer', 
                   'dog', 'frog', 'horse', 'ship', 'truck')
        
        print(f"训练集大小: {len(trainset)}")
        print(f"测试集大小: {len(testset)}")
        print(f"类别: {classes}")
        
        return trainloader, testloader, classes
    except Exception as e:
        print(f"数据加载失败: {e}")
        print("请确保已正确安装torchvision，或者尝试以下解决方案:")
        print("1. 重新安装torchvision: pip install --upgrade torchvision")
        print("2. 或者使用conda安装: conda install torchvision -c pytorch")
        sys.exit(1)

# 定义简单的卷积神经网络
class SimpleCIFAR10Net(nn.Module):
    def __init__(self, num_classes=10):
        super(SimpleCIFAR10Net, self).__init__()
        # 简化的网络结构
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(128 * 4 * 4, 256)
        self.fc2 = nn.Linear(256, num_classes)
        self.dropout = nn.Dropout(0.5)
        
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = self.pool(F.relu(self.conv4(x)))
        x = x.view(-1, 128 * 4 * 4)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 训练函数
def train_model(model, trainloader, criterion, optimizer, epoch):
    """训练模型"""
    model.train()  # 设置为训练模式
    running_loss = 0.0
    correct = 0
    total = 0
    
    print(f"Epoch {epoch+1} 开始训练...")
    start_time = time.time()
    
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        # 将数据移到设备
        inputs, targets = inputs.to(device), targets.to(device)
        
        # 梯度清零
        optimizer.zero_grad()
        
        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        
        # 反向传播
        loss.backward()
        
        # 参数更新
        optimizer.step()
        
        # 统计信息
        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        
        # 每50个batch打印一次信息
        if (batch_idx + 1) % 50 == 0:
            print(f'Epoch: {epoch+1} | Batch: {batch_idx+1} | Loss: {running_loss/(batch_idx+1):.3f} | Acc: {100.*correct/total:.3f}%')
    
    epoch_time = time.time() - start_time
    train_loss = running_loss / len(trainloader)
    train_acc = 100. * correct / total
    
    print(f'Epoch {epoch+1} 训练完成 | Loss: {train_loss:.3f} | Acc: {train_acc:.3f}% | Time: {epoch_time:.2f}s')
    
    return train_loss, train_acc

# 测试函数
def test_model(model, testloader, criterion):
    """测试模型"""
    model.eval()  # 设置为评估模式
    test_loss = 0
    correct = 0
    total = 0
    
    print("开始测试...")
    with torch.no_grad():
        for inputs, targets in testloader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
    
    test_loss /= len(testloader)
    test_acc = 100. * correct / total
    
    print(f'测试结果 | Loss: {test_loss:.3f} | Acc: {test_acc:.3f}%')
    
    return test_loss, test_acc

# 显示模型结构信息
def show_model_info(model):
    """显示模型信息"""
    print("\n模型结构信息:")
    print("-" * 50)
    
    total_params = 0
    for name, parameter in model.named_parameters():
        param_count = parameter.numel()
        total_params += param_count
        print(f"{name}: {list(parameter.shape)} = {param_count:,} 参数")
    
    print("-" * 50)
    print(f"总参数数量: {total_params:,}")
    print(f"可训练参数数量: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")

# 主函数
def main():
    """主函数"""
    print("PyTorch CIFAR10简单分类器")
    print("=" * 50)
    
    # 加载数据
    trainloader, testloader, classes = load_cifar10_data()
    
    # 创建模型
    model = SimpleCIFAR10Net(num_classes=10).to(device)
    
    # 显示模型信息
    show_model_info(model)
    
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    # 训练模型
    print("\n开始训练...")
    num_epochs = 5  # 减少epoch数量以加快演示
    train_losses = []
    train_accuracies = []
    test_losses = []
    test_accuracies = []
    
    start_time = time.time()
    
    for epoch in range(num_epochs):
        # 训练
        train_loss, train_acc = train_model(model, trainloader, criterion, optimizer, epoch)
        train_losses.append(train_loss)
        train_accuracies.append(train_acc)
        
        # 测试
        test_loss, test_acc = test_model(model, testloader, criterion)
        test_losses.append(test_loss)
        test_accuracies.append(test_acc)
        
        print("-" * 50)
    
    total_time = time.time() - start_time
    print(f"总训练时间: {total_time:.2f}秒")
    
    # 最终测试
    print("\n最终测试结果:")
    final_test_loss, final_test_acc = test_model(model, testloader, criterion)
    
    # 显示训练历史
    print("\n训练历史摘要:")
    print("Epoch\tTrain Loss\tTrain Acc\tTest Loss\tTest Acc")
    print("-" * 60)
    for i in range(num_epochs):
        print(f"{i+1}\t{train_losses[i]:.3f}\t\t{train_accuracies[i]:.3f}%\t\t{test_losses[i]:.3f}\t\t{test_accuracies[i]:.3f}%")
    
    # 保存模型
    torch.save(model.state_dict(), 'cifar10_simple_model.pth')
    print("\n模型已保存为 'cifar10_simple_model.pth'")
    
    print("\n程序执行完成!")

# 解释CIFAR10数据集
def explain_cifar10():
    """解释CIFAR10数据集"""
    print("\nCIFAR10数据集说明:")
    print("-" * 30)
    print("CIFAR10是一个常用的图像分类数据集，包含以下10个类别:")
    classes = ['飞机(airplane)', '汽车(automobile)', '鸟(bird)', '猫(cat)', '鹿(deer)', 
               '狗(dog)', '青蛙(frog)', '马(horse)', '船(ship)', '卡车(truck)']
    for i, cls in enumerate(classes):
        print(f"  {i}: {cls}")
    
    print("\n数据集规模:")
    print("  训练集: 50,000张 32x32彩色图像")
    print("  测试集: 10,000张 32x32彩色图像")
    print("  每个类别: 6,000张图像(训练集5,000 + 测试集1,000)")

# 解释神经网络训练过程
def explain_training_process():
    """解释神经网络训练过程"""
    print("\n神经网络训练过程说明:")
    print("-" * 30)
    print("1. 前向传播(Forward Propagation):")
    print("   - 输入数据通过网络逐层传递")
    print("   - 每层应用权重、偏置和激活函数")
    print("   - 最终得到预测输出")
    
    print("\n2. 损失计算(Loss Calculation):")
    print("   - 比较预测输出与真实标签")
    print("   - 使用交叉熵损失函数衡量差异")
    
    print("\n3. 反向传播(Backward Propagation):")
    print("   - 计算损失相对于每个参数的梯度")
    print("   - 使用链式法则从输出层向输入层传播")
    
    print("\n4. 参数更新(Parameters Update):")
    print("   - 使用优化器(Adam)更新网络权重")
    print("   - 沿梯度反方向调整参数以最小化损失")
    
    print("\n5. 评估(Evaluation):")
    print("   - 在测试集上评估模型性能")
    print("   - 计算准确率等指标")

if __name__ == '__main__':
    # 解释数据集
    explain_cifar10()
    
    # 解释训练过程
    explain_training_process()
    
    # 运行主程序
    main()